import json
from openai import OpenAI
from tqdm import tqdm  # 导入 tqdm 进度条库

# 配置 OpenAI 客户端
client = OpenAI(
    base_url='https://xiaoai.plus/v1',
    api_key='sk-7We2pzVKmzUpyP4kgD8gGAzRiBvQjUc7SV3oP8lLS9isoymj'  # 请替换为你的 API Key
)


# 调用大模型获取 JSON 格式的答案
def get_model_answer(question, options):
    """
    发送问题和选项给大模型，获取 JSON 格式的答案
    """
    prompt = (
             f" you will handle the following medical question. Please think carefully about each option "
            f"Question: {question}\n"
            f"Options:\n" + "\n".join([f"{key}. {value}" for key, value in options.items()]) + "\n"
                                                                                               f"Please provide your thought process and final answer in the following format:\n{{\"question\": \"Question content\", \"think\": \"Your thought process\", \"answer\": \"A\" }}. "
                                                                                               f"The answer field should only contain one of the letters: A, B, C, D, or E corresponding to the correct option."
    )


    while True:  # 进入循环，直到成功解析JSON
        try:
            completion = client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[{"role": "system", "content": "Please strictly adhere to the JSON format when answering."},
                          {"role": "user", "content": prompt}]
            )
            response = completion.choices[0].message.content.strip()

            print(response)
            # 提取 JSON 部分
            start_idx = response.find("{")
            end_idx = response.rfind("}") + 1
            json_str = response[start_idx:end_idx]  # 获取有效的 JSON 字符串

            # 尝试解析 JSON
            return json.loads(json_str)

        except Exception as e:
            print(f"JSON 解析错误: {e}")
            # 如果解析失败，重新请求大模型
            print("重新请求大模型...")
            continue  # 重新进入循环并重新发起请求

def get_processed_lines(output_file):
    try:
        with open(output_file, 'r', encoding='utf-8') as f_out:
            return sum(1 for _ in f_out)  # 统计已有的行数
    except FileNotFoundError:
        return 0  # 如果文件不存在，则从头开始处理

# 处理 JSONL 文件并即时保存模型返回的 JSON
def process_jsonl(input_file, output_file):
    """
    读取 JSONL 文件，从未处理的行开始调用大模型获取 JSON 结果，并存储到新 JSONL 文件（追加模式）。
    """
    processed_lines = get_processed_lines(output_file)  # 读取已处理的行数

    with open(input_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()  # 读取所有行
        total_lines = len(lines)  # 获取总行数

        if processed_lines >= total_lines:
            print("所有问题已处理完毕，无需重复处理。")
            return

        with tqdm(total=total_lines - processed_lines, desc="处理问题", unit="问题") as pbar, open(output_file, 'a', encoding='utf-8') as f_out:
            for line in lines[processed_lines:]:  # 从未处理的行开始
                data = json.loads(line)
                question = data['question']
                options = data['options']

                # 获取模型 JSON 结果
                model_response = get_model_answer(question, options)

                if model_response:
                    f_out.write(json.dumps(model_response, ensure_ascii=False) + "\n")  # 追加写入
                    f_out.flush()  # 及时刷新文件，防止数据丢失

                pbar.update(1)  # 更新进度条

    print(f"模型返回的 JSON 结果已追加保存至 {output_file}")



# 评估正确率函数
def evaluate_accuracy(input_file, output_file):
    """
    评估模型输出文件的正确率。
    """
    with open(input_file, 'r', encoding='utf-8') as f_in, open(output_file, 'r', encoding='utf-8') as f_out:
        input_data = [json.loads(line) for line in f_in]  # 读取输入文件中的所有数据
        output_data = [json.loads(line) for line in f_out]  # 读取输出文件中的所有数据

        correct = 0
        total = len(input_data)

        for input_item, output_item in zip(input_data, output_data):
            correct_answer = input_item['answer_idx']  # 从输入数据中获取正确答案的索引
            model_answer = output_item['answer']  # 获取模型预测的答案

            if str(correct_answer) == str(model_answer):  # 如果模型答案与正确答案相同
                correct += 1

        accuracy = (correct / total) * 100 if total > 0 else 0
        print(f"正确率: {accuracy:.2f}%")


# 主函数
def main():
    input_file = "G:\desktop\论文\data\med_qa\data_clean\questions\\US\结果\\test.jsonl"  # 你的输入 JSONL 文件路径
    output_file = "G:\desktop\论文\data\med_qa\data_clean\questions\\US\结果\\output_cot.jsonl"  # 结果保存路径
    process_jsonl(input_file, output_file)

    # 评估模型输出的准确率
    evaluate_accuracy(input_file, output_file)


# 运行主函数
if __name__ == "__main__":
    main()
